package com.ian.domain;

import com.ian.constant.Constants;
import com.ian.m3u8.exception.M3u8Exception;
import com.ian.m3u8.listener.DownloadListener;
import com.ian.m3u8.util.Constant;
import com.ian.enumeration.MediaFormat;
import com.ian.util.StringUtils;
import com.ian.util.TimeUtils;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.math.BigDecimal;
import java.net.HttpURLConnection;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Security;
import java.security.spec.AlgorithmParameterSpec;
import java.util.*;
import java.util.concurrent.*;

/**
 * @author Witt
 * @version 1.0.0
 * @desc M3u8Download 实体类
 * @date 2022/11/12 21:55
 */
@Data
@Slf4j
public class M3u8Download {

    /**
     * 解决java不支持AES/CBC/PKCS7Padding模式解密
     */
    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    // 优化内存占用
//    private final BlockingQueue<byte[]> BLOCKING_QUEUE = new LinkedBlockingQueue<>();

    // 要下载的m3u8链接
    private String downloadUrl;

    // 线程数
    private Integer threadCount = 1;

    // 重试次数
    private Integer retryCount = 30;

    // 链接连接超时时间（单位：毫秒）
    private long timeoutMillisecond = 2000L;

    // 合并后的文件存储目录
    private String dir;

    // 合并后的视频文件名称
    private String fileName;

    // 已完成ts片段个数
    private Integer finishedCount = 0;

    // 解密算法名称
    private String method;

    // 密钥
    private String key = "";

    // 密钥字节
    private byte[] keyBytes = new byte[16];

    // key是否为字节
    private boolean isByte = false;

    // IV
    private String iv = "";

    // 所有ts片段下载链接
    private Set<String> tsSet = new LinkedHashSet<>();

    // 所有ts片段下载链接
    private Map<String, Integer> tsMap = new HashMap<>();

    private static M3u8Download m3u8Download;

    // 解密后的片段
    private Set<File> finishedFiles = new ConcurrentSkipListSet<>(Comparator.comparingInt(o -> Integer.parseInt(o.getName().replace(".xyz", ""))));

    // 已经下载的文件大小
    private BigDecimal downloadBytes = new BigDecimal(0);

    // 监听间隔
    private volatile long interval = 0L;

    // 自定义请求头
    private Map<String, Object> requestHeaderMap = new HashMap<>();

    // 监听事件
    private Set<DownloadListener> listenerSet = new HashSet<>(5);

    // 代理设置
    private Proxy proxy;

    public void setThreadCount(Integer threadCount) {
        /*if (this.BLOCKING_QUEUE.size() < threadCount) {
            for (int i = this.BLOCKING_QUEUE.size(); i < threadCount * Constant.FACTOR; i++) {
                try {
                    this.BLOCKING_QUEUE.put(new byte[Constants.BYTE_COUNT]);
                } catch (InterruptedException ignored) {
                }
            }
        }*/
        this.threadCount = threadCount;
    }

    public void addRequestHeaderMap(M3u8Download m3u8Download, Map<String, Object> requestHeaderMap) {
        this.requestHeaderMap.putAll(requestHeaderMap);
    }

    public void addListener(DownloadListener downloadListener) {
        this.listenerSet.add(downloadListener);
    }

    public void setProxy(int port) {
        this.proxy = new Proxy(Proxy.Type.HTTP, new InetSocketAddress("127.0.0.1", port));
    }

    public void setProxy(String address, int port) {
        this.proxy = new Proxy(Proxy.Type.HTTP, new InetSocketAddress(address, port));
    }

    public void setProxy(Proxy.Type type, String address, int port) {
        this.proxy = new Proxy(type, new InetSocketAddress(address, port));
    }

    private M3u8Download(String downloadUrl) {
        this.downloadUrl = downloadUrl;
        requestHeaderMap.put("Referer", Constants.REFERER_LINUX);
        // requestHeaderMap.put("User-Agent", Constants.USER_AGENT_WINDOWS);
        requestHeaderMap.put("User-Agent", Constants.USER_AGENT_LINUX);
    }

    /**
     * 获取实例
     *
     * @param downloadUrl 要下载的链接
     * @return 返回m3u8下载实例
     */
    public static M3u8Download getInstance(String downloadUrl) {
        return new M3u8Download(downloadUrl);
    }


    /**
     * 开始下载视频
     */
    public void start() {
        long start = TimeUtils.getTime(Constants.TIME_SECOND);
        checkField();
        String tsUrl = getTsUrl();
        if (StringUtils.isBlank(tsUrl)) {
            log.info("不需要解密");
        }
        startDownload();
        log.info("m3u8 视频下载任务执行完毕，用时 {} 秒！", (TimeUtils.getTime(Constants.TIME_SECOND) - start));
    }


    /**
     * 下载视频
     */
    private void startDownload() {
        // 如果生成目录不存在，则创建
        File baseDir = new File(dir);
        if (!baseDir.exists()) {
            baseDir.mkdirs();
        }

        int tsMapSize = tsMap.entrySet().size();
        ForkJoinPool.commonPool().getParallelism();
        /*
        tsMap.entrySet().stream().parallel().forEach(entry -> {...}
        tsMap.entrySet().parallelStream().forEach(entry -> {...}
        上面2种并行流的执行，都可以搭配下面的并行线程数量的设置。实际测试是设置为100最快（针对我的电脑上的这个项目）。
        System.setProperty("java.util.concurrent.ForkJoinPool.common.parallelism", "100")
         */
//        tsMap.entrySet().stream().parallel().forEach(entry -> {
        tsMap.entrySet().parallelStream().forEach(entry -> {
            String tsUrl = entry.getKey();
            Integer cnt = entry.getValue();
            int count = 1;
            HttpURLConnection httpURLConnection = null;
            // xy为未解密的ts片段，如果存在，则删除
            File originalFile = new File(dir + Constants.FILESEPARATOR + cnt + ".xy");
            if (originalFile.exists()) {
                originalFile.delete();
            }
            try {
                originalFile.createNewFile();
            } catch (IOException e) {
                e.printStackTrace();
            }

            File decryptedFile = new File(dir + Constants.FILESEPARATOR + cnt + ".xyz");
            if (decryptedFile.exists()) {
                decryptedFile.delete();
            }
            try {
                decryptedFile.createNewFile();
            } catch (IOException e) {
                e.printStackTrace();
            }

            byte[] bytes = new byte[Constants.BYTE_COUNT];
            /*
            try {
                // 此处可能有多个线程获取到同一个字节数组对象的bug，因为 BLOCKING_QUEUE 对象是所有线程共享的
                bytes = BLOCKING_QUEUE.take();
            } catch (InterruptedException e) {
                bytes = new byte[Constants.BYTE_COUNT];
            }*/
            // 重试次数判断
            while (count <= retryCount) {
                try (InputStream originalIS = new FileInputStream(originalFile);
                     OutputStream originalOS = new FileOutputStream(originalFile);
                     FileOutputStream decryptedOS = new FileOutputStream(decryptedFile);
                     ) {
                    // 模拟http请求获取ts片段文件
                    URL url = new URL(tsUrl);
                    if (proxy == null) {
                        httpURLConnection = (HttpURLConnection) url.openConnection();
                    } else {
                        httpURLConnection = (HttpURLConnection) url.openConnection(proxy);
                    }
                    httpURLConnection.setConnectTimeout((int) timeoutMillisecond);
                    for (Map.Entry<String, Object> entry2 : requestHeaderMap.entrySet()) {
                        httpURLConnection.addRequestProperty(entry2.getKey(), entry2.getValue().toString());
                    }
                    httpURLConnection.setUseCaches(false);
                    httpURLConnection.setReadTimeout((int) timeoutMillisecond);
                    httpURLConnection.setDoInput(true);
                    InputStream httpURLIS = httpURLConnection.getInputStream();

                    // 将未解密的ts片段写入文件
                    int len;
                    int available = 0;
                    // Thread.sleep(1000L);
                    while (available == 0) {
                        available = httpURLIS.available();
                    }
                    if (bytes.length < available) {
                        bytes = new byte[available];
                    }
                    while ((len = httpURLIS.read(bytes)) != -1) {
                        originalOS.write(bytes, 0, len);
                        synchronized (this) {
                            downloadBytes = downloadBytes.add(new BigDecimal(len));
                        }
                    }
                    originalOS.flush();
                    httpURLIS.close();
                    available = originalIS.available();
                    if (bytes.length < available) {
                        bytes = new byte[available];
                    }
                    originalIS.read(bytes, 0, available);
                    // 开始解密ts片段，把解密后的ts片段文件的后缀改为了xyz
                    //   判断Key是否存在、加解密类型是否正确
                    if (StringUtils.isBlank(key) || "NONE".equals(method)) {
                        decryptedOS.write(bytes, 0, available);
                    } else {
                        byte[] decrypt = decrypt(bytes, available, key, iv, method);
                        if (decrypt == null) {
                            decryptedOS.write(bytes, 0, available);
                        } else {
                            decryptedOS.write(decrypt);
                        }
                    }
                    // finishedFiles.add(decryptedFile);
                    finishedFiles.add(decryptedFile);
                    break;
                } catch (Exception e) {
                    if (e instanceof InvalidKeyException || e instanceof InvalidAlgorithmParameterException) {
                        log.error("解密失败！");
                        break;
                    } else {
                        log.debug("连接出错：{}", e.toString());
                    }
                    log.debug("第 {} 次获取链接重试：{}", count, tsUrl);
                    count++;
                } finally {
                    /*try {
                        BLOCKING_QUEUE.put(bytes);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }*/
                    if (httpURLConnection != null) {
                        httpURLConnection.disconnect();
                    }
                }
            }
            if (count > retryCount) {
                // 自定义异常
                //  throw new M3u8Exception("连接超时！");
                log.error("超出重试次数，连接出错！");
            }
            finishedCount++;
            log.info(tsUrl + " 下载完毕！\t已完成" + finishedCount + "个，还剩" + (tsMapSize - finishedCount) + "个！");
        });

        log.info("下载完成，正在合并文件！共" + finishedFiles.size() + "个！" + StringUtils.convertToDownloadSpeed(downloadBytes, 3));
        //开始合并视频
        mergeTs();
        // 清理临时内容：将所有在工作目录中的视频文件移动到上一层目录，然后删除工作目录以及其中的所有内容
        clearTmpContent();
        log.info("视频合并完成，欢迎使用!");
    }

    /**
     * 下载并合并m3u8链接的视频文件后，清理临时内容：将所有在工作目录中的视频文件移动到上一层目录，然后删除工作目录以及其中的所有内容
      */
    private void clearTmpContent() {
        try {
            Path workDir = Paths.get(this.dir);
            Files.walk(workDir)
                    .filter(path -> !workDir.equals(path))
                    .forEach(path -> {
                        // 删除下载好的片段
                        if (Files.isRegularFile(path) && path.toString().endsWith(".mp4")) {
                            try {
                                log.info("moving file '{}' to its parent folder...", path);
                                Path targetPath = path.getParent().getParent().resolve(path.getFileName());
                                // 将子目录中的视频文件移动到上一层目录
                                Files.move(path, targetPath);
                                log.info("将所有在子目录中的视频文件移动到上一层目录，成功！");
                            } catch (IOException e) {
                                log.error("将所有在子目录中的视频文件移动到上一层目录，然后删除子目录，失败！\n失败原因：{}", e.toString());
                            }
                        } else {
                            try {
                                Files.delete(path);
                            } catch (IOException e) {
                                throw new RuntimeException(e);
                            }
                        }
                    });
            // 删除下载目录
            Files.delete(workDir);
        } catch (IOException e) {
            log.error(e.toString());
        }
    }

    private void startListener(ExecutorService fixedThreadPool) {
        new Thread(() -> {
            for (DownloadListener downloadListener : listenerSet) {
                downloadListener.start();
            }
            // 轮询是否下载成功
            while (!fixedThreadPool.isTerminated()) {
                try {
                    Thread.sleep(interval);
                    for (DownloadListener downloadListener : listenerSet) {
                        downloadListener.process(downloadUrl, finishedCount, tsSet.size(), new BigDecimal(finishedCount).divide(new BigDecimal(tsSet.size()), 4, BigDecimal.ROUND_HALF_UP).multiply(new BigDecimal(100)).setScale(2, BigDecimal.ROUND_HALF_UP).floatValue());
                    }
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
            for (DownloadListener downloadListener : listenerSet) {
                downloadListener.end();
            }
        }).start();
        new Thread(() -> {
            while (!fixedThreadPool.isTerminated()) {
                try {
                    BigDecimal bigDecimal = new BigDecimal(downloadBytes.toString());
                    Thread.sleep(1000L);
                    for (DownloadListener downloadListener : listenerSet) {
                        downloadListener.speed(StringUtils.convertToDownloadSpeed(new BigDecimal(downloadBytes.toString()).subtract(bigDecimal), 3) + "/s");
                    }
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }).start();
    }

    /**
     * 合并下载好的ts片段
     */
    private void mergeTs() {
        FileOutputStream fileOutputStream = null;
        try {
            File file = new File(dir + Constants.FILESEPARATOR + fileName + ".mp4");
            System.gc();
            if (file.exists()) {
                file.delete();
            } else {
                file.createNewFile();
            }
            fileOutputStream = new FileOutputStream(file);
            byte[] b = new byte[40960];
            for (File f : finishedFiles) {
                FileInputStream fileInputStream = new FileInputStream(f);
                int len;
                int available = fileInputStream.available();
                if (b.length < available) {
                    b = new byte[available];
                }
                while ((len = fileInputStream.read(b)) != -1) {
                    fileOutputStream.write(b, 0, len);
                }
                fileInputStream.close();
                fileOutputStream.flush();
            }
            fileOutputStream.close();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (fileOutputStream != null) {
                try {
                    fileOutputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 获取所有的ts片段下载链接
     *
     * @return 链接是否被加密，null为非加密
     */
    private String getTsUrl() {
        StringBuilder content = getUrlContent(downloadUrl, false);
        // 判断是否是m3u8链接
        if (!content.toString().contains("#EXTM3U")) {
            throw new M3u8Exception(downloadUrl + "不是m3u8链接！");
        }
        String[] split = content.toString().split("\\n");
        String keyUrl = "";
        boolean isKey = false;
        for (String s : split) {
            // 如果含有此字段，则说明只有一层m3u8链接
            if (s.contains("#EXT-X-KEY") || s.contains("#EXTINF")) {
                isKey = true;
                keyUrl = downloadUrl;
                break;
            }
            // 如果含有此字段，则说明ts片段链接需要从第二个m3u8链接获取
            if (s.contains(".m3u8")) {
                if (StringUtils.isUrl(s)) {
                    return s;
                }
                String relativeUrl = downloadUrl.substring(0, downloadUrl.lastIndexOf("/") + 1);
                if (s.startsWith("/")) {
                    s = s.replaceFirst("/", "");
                }
                keyUrl = mergeUrl(relativeUrl, s);
                break;
            }
        }
        if (StringUtils.isBlank(keyUrl)) {
            throw new M3u8Exception("未发现key链接！");
        }
        // 获取密钥
        String key1 = isKey ? getKey(keyUrl, content) : getKey(keyUrl, null);
        if (StringUtils.isNotBlank(key1)) {
            key = key1;
        } else {
            key = null;
        }
        return key;
    }

    /**
     * 获取ts解密的密钥，并把ts片段加入set集合
     *
     * @param url     密钥链接，如果无密钥的m3u8，则此字段可为空
     * @param content 内容，如果有密钥，则此字段可以为空
     * @return ts是否需要解密，null为不解密
     */
    private String getKey(String url, StringBuilder content) {
        StringBuilder urlContent;
        if (StringUtils.isBlank(content)) {
            urlContent = getUrlContent(url, false);
        } else {
            urlContent = content;
        }
        if (!urlContent.toString().contains("#EXTM3U")) {
            throw new M3u8Exception(downloadUrl + "不是m3u8链接！");
        }
        String[] split = urlContent.toString().split("\\n");
        for (String s : split) {
            // 如果含有此字段，则获取加密算法以及获取密钥的链接
            if (s.contains("EXT-X-KEY")) {
                String[] split1 = s.split(",");
                for (String s1 : split1) {
                    if (s1.contains("METHOD")) {
                        method = s1.split("=", 2)[1];
                        continue;
                    }
                    if (s1.contains("URI")) {
                        key = s1.split("=", 2)[1];
                        continue;
                    }
                    if (s1.contains("IV")) {
                        iv = s1.split("=", 2)[1];
                    }
                }
            }
        }
        String relativeUrl = url.substring(0, url.lastIndexOf("/") + 1);
        /* 将ts片段链接加入set集合
        for (int i = 0; i < split.length; i++) {
            String s = split[i];
            if (s.contains("#EXTINF")) {
                String s1 = split[++i];
                tsSet.add(StringUtils.isUrl(s1) ? s1 : mergeUrl(relativeUrl, s1));
            }
        }
        */

        // 将ts片段链接放入 HashMap
        Integer cnt = 0;
        for (int i = 0; i < split.length; i++) {
            String s = split[i];
            if (s.contains("#EXTINF")) {
                String s1 = split[++i];
                Map<Integer, String> map = new HashMap<>();
                String tsUrl = StringUtils.isUrl(s1) ? s1 : mergeUrl(relativeUrl, s1);
                if (!tsMap.containsKey(tsUrl)) {
                    tsMap.put(tsUrl, ++cnt);
                }
                log.info("map内容：key：{}，value：{}", tsUrl, cnt);
            }
        }

        if (StringUtils.isNotBlank(key)) {
            key = key.replace("\"", "");
            return getUrlContent(StringUtils.isUrl(key) ? key : mergeUrl(relativeUrl, key), true).toString().replaceAll("\\s+", "");
        }
        return null;
    }

    /**
     * 模拟http请求获取内容
     *
     * @param urls  http链接
     * @param isKey 这个url链接是否用于获取key
     * @return 内容
     */
    private StringBuilder getUrlContent(String urls, boolean isKey) {
        int count = 1;
        HttpURLConnection httpURLConnection = null;
        StringBuilder content = new StringBuilder();
        while (count <= retryCount) {
            try {
                URL url = new URL(urls);
                if (proxy == null) {
                    httpURLConnection = (HttpURLConnection) url.openConnection();
                } else {
                    httpURLConnection = (HttpURLConnection) url.openConnection(proxy);
                }
                httpURLConnection.setConnectTimeout((int) timeoutMillisecond);
                httpURLConnection.setReadTimeout((int) timeoutMillisecond);
                httpURLConnection.setUseCaches(false);
                httpURLConnection.setDoInput(true);
                for (Map.Entry<String, Object> entry : requestHeaderMap.entrySet()) {
                    httpURLConnection.addRequestProperty(entry.getKey(), entry.getValue().toString());
                }
                String line;
                InputStream inputStream = httpURLConnection.getInputStream();
                BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(inputStream));
                if (isKey) {
                    byte[] bytes = new byte[128];
                    int len;
                    len = inputStream.read(bytes);
                    isByte = true;
                    if (len == 1 << 4) {
                        keyBytes = Arrays.copyOf(bytes, 16);
                        content.append("isByte");
                    } else {
                        content.append(new String(Arrays.copyOf(bytes, len)));
                    }
                    return content;
                }
                while ((line = bufferedReader.readLine()) != null) {
                    content.append(line).append("\n");
                }
                bufferedReader.close();
                inputStream.close();
                log.info(content.toString());
                break;
            } catch (Exception e) {
                log.debug("第" + count + "次获取链接重试！\t" + urls);
                count++;
                e.printStackTrace();
            } finally {
                if (httpURLConnection != null) {
                    httpURLConnection.disconnect();
                }
            }
        }
        if (count > retryCount) {
            throw new M3u8Exception("连接超时！");
        }
        return content;
    }

    /**
     * 解密ts
     *
     * @param sSrc   ts文件字节数组
     * @param length
     * @param sKey   密钥
     * @return 解密后的字节数组
     */
    private byte[] decrypt(byte[] sSrc, int length, String sKey, String iv, String method) throws Exception {
        // 判断Key是否存在、加解密类型是否正确
        if (StringUtils.isBlank(sKey) || "NONE".equals(method)) {
            return null;
        }

        if (StringUtils.isNotBlank(method) && !method.contains("AES")) {
            throw new M3u8Exception("未知的算法！");
        }

        // 判断Key是否为16位
        if (sKey.length() != 16 && !isByte) {
            throw new M3u8Exception("Key长度不是16位！");
        }
        Cipher cipher = Cipher.getInstance("AES/CBC/PKCS7Padding");
        SecretKeySpec keySpec = new SecretKeySpec(isByte ? keyBytes : sKey.getBytes(StandardCharsets.UTF_8), "AES");
        byte[] ivByte;
        if (iv.startsWith("0x")) {
            ivByte = StringUtils.hexStringToByteArray(iv.substring(2));
        } else {
            ivByte = iv.getBytes();
        }
        if (ivByte.length != 16) {
            ivByte = new byte[16];
        }
        // 如果m3u8有IV标签，那么IvParameterSpec构造函数就把IV标签后的内容转成字节数组传进去
        AlgorithmParameterSpec paramSpec = new IvParameterSpec(ivByte);
        cipher.init(Cipher.DECRYPT_MODE, keySpec, paramSpec);
        return cipher.doFinal(sSrc, 0, length);
    }

    /**
     * 字段校验
     */
    private void checkField() {
        if (!"m3u8".equals(MediaFormat.getMediaFormat(this.downloadUrl))) {
            throw new M3u8Exception(this.downloadUrl + "不是一个完整m3u8链接！");
        }
        if (this.threadCount <= 0) {
            throw new M3u8Exception("同时下载线程数只能大于0！");
        }
        if (this.retryCount < 0) {
            throw new M3u8Exception("重试次数不能小于0！");
        }
        if (this.timeoutMillisecond < 0) {
            throw new M3u8Exception("超时时间不能小于0！");
        }
        if (StringUtils.isBlank(this.dir)) {
            throw new M3u8Exception("视频存储目录不能为空！");
        }
        if (StringUtils.isBlank(this.fileName)) {
            throw new M3u8Exception("视频名称不能为空！");
        }
        this.finishedCount = 0;
        this.method = "";
        this.key = "";
        this.isByte = false;
        this.iv = "";
        this.tsSet.clear();
        this.finishedFiles.clear();
        this.downloadBytes = new BigDecimal(0);
    }

    private String mergeUrl(String start, String end) {
        if (end.startsWith("/")) {
            end = end.replaceFirst("/", "");
        }
        int position = 0;
        String subEnd, tempEnd = end;
        while ((position = end.indexOf("/", position)) != -1) {
            subEnd = end.substring(0, position + 1);
            if (start.endsWith(subEnd)) {
                tempEnd = end.replaceFirst(subEnd, "");
                break;
            }
            ++position;
        }
        return start + tempEnd;
    }


}
